

package com.mantou.mantouaiagent.advisor;

import com.mantou.mantouaiagent.util.ThreadLocalUtil;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.MessageAggregator;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

/**
 * A simple logger advisor that logs the request and response messages.
 *
 * @author Christian Tzolov
 */
@Component
public class MyLoggerAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

	private static final Logger log = LogManager.getLogger(MyLoggerAdvisor.class);

	@Override
	public String getName() {
		return this.getClass().getSimpleName();
	}

	@Override
	public int getOrder() {
		return 0;
	}

	private AdvisedRequest before(AdvisedRequest request) {
		log.info("AI request: {}", request.userText());
//		AdvisedRequest advisedRequest = AdvisedRequest.from(request).systemText("提示词").build();
		return request;
	}

	private void observeAfter(AdvisedResponse advisedResponse) {
		log.info("Ai response: {}", advisedResponse.response().getResult().getOutput().getText());
	}

	@Override
	public String toString() {
		return MyLoggerAdvisor.class.getSimpleName();
	}

	@Override
	public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
//		AdvisedRequest request = AdvisedRequest.from(advisedRequest).systemText("提示词").build();
		advisedRequest = before(advisedRequest);


		AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);

		observeAfter(advisedResponse);

		return advisedResponse;
	}

	@Override
	public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {

		advisedRequest = before(advisedRequest);

		Flux<AdvisedResponse> advisedResponses = chain.nextAroundStream(advisedRequest);

		return new MessageAggregator().aggregateAdvisedResponse(advisedResponses, this::observeAfter);
	}

}
